package org.annoflow.javassist;

import java.security.AccessController;
import java.security.PrivilegedActionException;
import java.security.PrivilegedExceptionAction;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import javassist.CannotCompileException;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.NotFoundException;
import javassist.expr.Cast;
import javassist.expr.ConstructorCall;
import javassist.expr.Expr;
import javassist.expr.ExprEditor;
import javassist.expr.FieldAccess;
import javassist.expr.Handler;
import javassist.expr.Instanceof;
import javassist.expr.MethodCall;
import javassist.expr.NewArray;
import javassist.expr.NewExpr;


import org.apache.log4j.Logger;
import org.annoflow.filter.FilterPoint;
import org.annoflow.javassist.AbstractCommand.SignatureType;
import org.annoflow.policy.PolicyType;
import org.annoflow.policy.Policy;

/**
 * The Class CustomExprEditor. This class is responsible for actually editing
 * the
 */
public class AnnoCustomExprEditor extends ExprEditor {
	protected Logger logger = Logger.getLogger(CustomExprEditor.class);

	private abstract class PrivlegedJavassistAction implements
			PrivilegedExceptionAction<String> {
		protected Object context;

		public void setContext(Object context) {
			this.context = context;
		}
	}

	Map<SignatureType, Map<String, List<CommandInterface>>> ExpressionNameToCommand;

	/**
	 * Instantiates a new custom expr editor. from a list of commands. These
	 * commands are stored in seperate lists to be executed on method calls and
	 * expressions being translated.
	 * 
	 * @param editConfigs
	 *            commands to drive the configuration.
	 */
	public AnnoCustomExprEditor(List<CommandInterface> editConfigs) {
		ExpressionNameToCommand = new HashMap<SignatureType, Map<String, List<CommandInterface>>>();
		for (CommandInterface curCommand : editConfigs) {

			// CommandInterface curCommand = commandIter.next();
			Map<String, List<CommandInterface>> sigMap = ExpressionNameToCommand
					.get(curCommand.getSignatureType());
			if (sigMap == null) {
				sigMap = new HashMap<String, List<CommandInterface>>();
				ExpressionNameToCommand.put(curCommand.getSignatureType(),
						sigMap);
			}
			List<CommandInterface> commands = sigMap.get(curCommand
					.getSignature());
			if (commands == null) {
				commands = new LinkedList<CommandInterface>();
				sigMap.put(curCommand.getSignature(), commands);
			}
			commands.add(curCommand);
		}
	}

	/**
	 * this method is for static edits to a class. These happen once and are not
	 * iterated.
	 * 
	 * @param clazz
	 *            the clazz to be edited
	 */
	public synchronized void staticEdits(CtClass clazz) {
		Map<String, List<CommandInterface>> sigTypeMap = ExpressionNameToCommand
				.get(SignatureType.METHODBODY);
		performStaticEdits(clazz, sigTypeMap);
		sigTypeMap = ExpressionNameToCommand
				.get(SignatureType.STATICDEFINITION);
		performStaticEdits(clazz, sigTypeMap);
	}

	private void performStaticEdits(CtClass clazz,
			Map<String, List<CommandInterface>> sigTypeMap) {
		if (sigTypeMap != null) {
			for (CtMethod method : clazz.getDeclaredMethods()) {
				String methodName = method.getLongName();
				int nameStart = method.getLongName()
						.substring(0, method.getLongName().indexOf('('))
						.lastIndexOf('.');
				methodName = methodName.substring(nameStart + 1);
				List<CommandInterface> listOfCommands = sigTypeMap
						.get(methodName);
				if (listOfCommands != null) {
					for (CommandInterface curCommand : listOfCommands) {
						// CommandInterface curCommand = iterCommand.next();
						try {
							curCommand.translate(method, methodName);
						} catch (Exception e) {
							logger.error(
									"Could not perform Static edit: ["
											+ methodName + ","
											+ curCommand.getSignature() + "]",
									e);
						}
					}
				}
			}
		}
	}

	/**
	 * Conversion of an Expr.
	 * 
	 * @param ccall
	 *            the expression ccall
	 * @param commandList
	 *            the list of commands to use to perform the edit.
	 * @param expressionName
	 *            the name of the expression being edited.
	 */
	private void convert(Expr ccall, List<CommandInterface> commandList,
			boolean isSuper) {
		if (commandList != null) {
			for (CommandInterface curCommand : commandList) {

				try {
					curCommand.translate(ccall, isSuper);
				} catch (MalformedCommandException e) {
					logger.error(
							"The command is not written correctly please revise it:  "
									+ curCommand.toString(), e);
				} catch (CannotCompileException e) {
					logger.error("Could not compile the replacement text: "
							+ curCommand.getReplacement(), e);
				}
			}
		

		}
		else
		{
			if(ccall instanceof MethodCall && !(ccall instanceof ConstructorCall))
			{
				try {
				boolean rewrite = true;
				MethodCall call = (MethodCall)ccall;
				try
				{
					call.getMethod().getAnnotation(Policy.class);
					
				}catch (ClassNotFoundException e) {
					rewrite = false;
				}
				CtClass clazz = call.getMethod().getDeclaringClass();
				if(clazz.getName().startsWith("org.annoflow") || clazz.getName().startsWith("javassist") || clazz.getName().startsWith("java"))
				{
					rewrite = false;
				}
				
				if(rewrite)
						call.replace("org.annoflow.audit.ContextTracker.getInstance().setCallerContext(\""+ call.getEnclosingClass().getName()+"\");$_=$proceed($$);org.annoflow.audit.ContextTracker.getInstance().clear();");
			
			} catch (CannotCompileException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			} catch (NotFoundException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			} 

		}
		}

	}

	public List<CommandInterface> getCommandList(SignatureType sigType,
			String signature) {
		Map<String, List<CommandInterface>> callToCommand = ExpressionNameToCommand
				.get(sigType);
		if (callToCommand != null) {
			return callToCommand.get(signature);
		}
		return null;
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see javassist.expr.ExprEditor#edit(javassist.expr.Handler)
	 */
	public void edit(Handler ccall) {
		/*
		 * try { PrivlegedJavassistAction action = new
		 * PrivlegedJavassistAction() {
		 * 
		 * @Override public String run() throws NotFoundException { CtClass type
		 * = ((Handler)context).getType(); if(type ==null) { return ""; } return
		 * type.getName(); }
		 * 
		 * }; action.setContext(ccall); String constName =
		 * AccessController.doPrivileged(action);
		 * logger.trace("Attempting to rewrite "+constName); //String handlerSig
		 * = ccall.getType().getName(); convert(ccall,
		 * getCommandList(SignatureType.HANDLER, constName), false); } catch
		 * (PrivilegedActionException e) {
		 * logger.error("Exception occured rewriting a Handler call in "
		 * +ccall.getFileName()); e.printStackTrace(); }
		 */

	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see javassist.expr.ExprEditor#edit(javassist.expr.Cast)
	 */
	public void edit(Cast ccall) {

		/*
		 * try { PrivlegedJavassistAction action = new
		 * PrivlegedJavassistAction() {
		 * 
		 * @Override public String run() throws NotFoundException { return
		 * ((Cast)context).getType().getName(); }
		 * 
		 * }; action.setContext(ccall); String constName =
		 * AccessController.doPrivileged(action);
		 * logger.trace("Attempting to rewrite "+constName); convert(ccall,
		 * getCommandList(SignatureType.CAST, constName), false); } catch
		 * (PrivilegedActionException e) {
		 * logger.error("Exception occured rewriting a Cast call in "
		 * +ccall.getFileName()); e.printStackTrace(); }
		 */
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see javassist.expr.ExprEditor#edit(javassist.expr.Instanceof)
	 */
	public void edit(Instanceof ccall) {
		/*
		 * try { PrivlegedJavassistAction action = new
		 * PrivlegedJavassistAction() {
		 * 
		 * @Override public String run() throws NotFoundException { return
		 * ((Instanceof)context).getType().getName(); }
		 * 
		 * }; action.setContext(ccall); String constName =
		 * AccessController.doPrivileged(action);
		 * logger.trace("Attempting to rewrite "+constName); convert(ccall,
		 * getCommandList(SignatureType.INSTANCEOF, constName), false); } catch
		 * (PrivilegedActionException e) {
		 * logger.error("Exception occured rewriting a Instance of call in "
		 * +ccall.getFileName()); e.printStackTrace(); }
		 */

	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see javassist.expr.ExprEditor#edit(javassist.expr.FieldAccess)
	 */
	public void edit(FieldAccess ccall) {
		try {
			Object[] annotations = ccall.getField().getAnnotations();
			for (Object anno : annotations) {
				if (anno instanceof Policy) {
					Policy annoPolicy = (Policy) anno;
					Class<? extends PolicyType> policyClass = annoPolicy
							.policy();
					PolicyType policy = policyClass.newInstance();
					if (ccall.isWriter()) {
						logger.info("Adding policy to write op on field "
								+ ccall.getFieldName() + " in class "
								+ ccall.getEnclosingClass().getName());
						ccall.replace("$_=$proceed($$);org.annoflow.policy.PolicyManager.addFieldPolicy("
								+ ccall.getFieldName()
								+ ","
								+ "this"
								+ ","
								+ policy.getClass().getCanonicalName()
								+ ".class);");
					}

				}
			}
		} catch (ClassNotFoundException e) {
			e.printStackTrace();
		} catch (NotFoundException e) {
			e.printStackTrace();
		} catch (InstantiationException e) {
			e.printStackTrace();
		} catch (IllegalAccessException e) {
			e.printStackTrace();
		} catch (CannotCompileException e) {
			e.printStackTrace();
		}

		// TODO: This probably does not work! If you attempt to access the field
		// in a super class this may fall apart.
		/*
		 * try { action.setClass(ccall.getType()); String fieldName =
		 * ccall.getField().getName(); convert(ccall,
		 * getCommandList(SignatureType.FIELDACCESS, fieldName), false); } catch
		 * (NotFoundException e) {
		 * System.err.println("Could not find class for array type");
		 * e.printStackTrace(); }
		 */
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see javassist.expr.ExprEditor#edit(javassist.expr.NewArray)
	 */
	public void edit(NewArray ccall) {

		/*
		 * try { PrivlegedJavassistAction action = new
		 * PrivlegedJavassistAction() {
		 * 
		 * @Override public String run() throws NotFoundException { return
		 * ((NewArray)context).getComponentType().getName(); }
		 * 
		 * }; action.setContext(ccall); String constName =
		 * AccessController.doPrivileged(action);
		 * logger.trace("Attempting to rewrite "+constName); convert(ccall,
		 * getCommandList(SignatureType.NEWARRAY, constName), false); } catch
		 * (PrivilegedActionException e) {
		 * logger.error("Exception occured rewriting a NewArray call in "
		 * +ccall.getFileName()); e.printStackTrace(); }
		 */

	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see javassist.expr.ExprEditor#edit(javassist.expr.MethodCall)
	 */
	public void edit(MethodCall ccall) {
		try {
			PrivlegedJavassistAction action = new PrivlegedJavassistAction() {

				@Override
				public String run() throws NotFoundException {
					return ((MethodCall) context).getMethod().getLongName();
				}

			};
			action.setContext(ccall);
			String constName = (String) AccessController.doPrivileged(action);
			logger.trace("Attempting to rewrite " + constName);
			convert(ccall, getCommandList(SignatureType.METHODCALL, constName),
					ccall.isSuper());
		} catch (PrivilegedActionException e) {
			logger.error("Exception occured rewriting a MethodCall call in "
					+ ccall.getFileName(), e);
		}
		
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see javassist.expr.ExprEditor#edit(javassist.expr.ConstructorCall)
	 */
	public void edit(ConstructorCall ccall) {
		try {
			PrivlegedJavassistAction action = new PrivlegedJavassistAction() {

				@Override
				public String run() throws NotFoundException {
					return ((ConstructorCall) context).getConstructor()
							.getLongName();
				}

			};
			action.setContext(ccall);
			String constName = AccessController.doPrivileged(action);
			logger.trace("Attempting to rewrite " + constName);
			convert(ccall,
					getCommandList(SignatureType.CONSTRUCTORCALL, constName),
					ccall.isSuper());
		} catch (PrivilegedActionException e) {
			logger.error(
					"Exception occured rewriting a ConstructorCall call in "
							+ ccall.getFileName(), e);

		}

	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see javassist.expr.ExprEditor#edit(javassist.expr.NewExpr)
	 */
	public void edit(NewExpr ccall) {
		try {
			PrivlegedJavassistAction action = new PrivlegedJavassistAction() {

				@Override
				public String run() throws NotFoundException {
					return ((NewExpr) context).getConstructor().getLongName();
				}

			};
			action.setContext(ccall);
			String constName = AccessController.doPrivileged(action);
			logger.trace("Attempting to rewrite " + constName);
			List<CommandInterface> commands = getCommandList(
					SignatureType.NEWEXPR, constName);
			if (commands != null) {
				convert(ccall, commands, false);
			} else if(ccall.getEnclosingClass().hasAnnotation(Policy.class)){
				checkPolicy(ccall);
			}
		} catch (PrivilegedActionException e) {
			logger.error("Exception occured rewriting a NewExpr call in "
					+ ccall.getFileName(), e);

		} catch (CannotCompileException e2) {
			logger.error(
					"Exception occured when adding the policy check to the new Expression",
					e2);
		}

	}

	private void checkPolicy(NewExpr expr) throws CannotCompileException {
		expr.replace("$_=$proceed($$);org.annoflow.audit.ContextTracker.getInstance().setCallerContext(\""+expr.getEnclosingClass().getName()+"\");org.annoflow.policy.PolicyManager.addObjectPolicy($_);org.annoflow.audit.ContextTracker.getInstance().clear();");
	}

}